from langgraph.graph import StateGraph, MessagesState
from langgraph.checkpoint.memory import InMemorySaver
from langgraph_memory.src.node.node import update_memory, call_model
from langgraph_memory.src.memory.memory_store import RedisMilvusStore, init_milvus_collection
from langgraph_memory.src.memory.embedding import EmbeddingModel
from redis import Redis
from langgraph_memory.src.config.config import Settings

# 初始化组件
embedder = EmbeddingModel()
milvus_collection = init_milvus_collection(embedder)
redis_client = Redis(host=Settings.REDIS_HOST, port=Settings.REDIS_PORT, decode_responses=True)
custom_store = RedisMilvusStore(redis_client, milvus_collection, embedder)
checkpointer = InMemorySaver()

# 构建图
builder = StateGraph(MessagesState)

builder.add_node("update_memory", update_memory)
builder.add_node("call_model", call_model)

builder.set_entry_point("update_memory")
builder.add_edge("update_memory", "call_model")
builder.add_edge("call_model", "__end__")

# 编译图
graph = builder.compile(checkpointer=checkpointer, store=custom_store)